#include "RenderState.h"
#include "../Misc/Crc.h"
#include "../Templates/NewObject.h"
#include "../Core/CoreGlobals.h"
#include "../Core/MaggotApp.h"
#include "Swapchain.h"
#include "../Game/GameState.h"
#include "../Thread/ThreadManager.h"

#include "Texture/Picture2DTexture.h"
#include "Imgui/ImguiGUI.h"

RenderState::RenderState()
{
	mFrameIndex = 0;
	mFrameAmount = 0;
	mFrameCount = 0;
	mIsSubmitTexture = false;
	mImmediateCommandBuffer = VK_NULL_HANDLE;
}

RenderState::~RenderState()
{
	for (uint32 i = 0; i < mFrameCount; ++i)
	{
		CloseHandle(mFrameEvents[i]);
	}
	mFrameEvents.clear();
}

void RenderState::Init()
{
	mRenderProcessMap = MaggotApp::Get().CreateRenderProcess();

	CreateCommandPool();

	CreateStateResource();

	CreateSwapchain();

	mRenderTimer.Reset();

	// create event for frame
	mFrameEvents.resize(mFrameCount);
	for (uint32 i = 0; i < mFrameCount; ++i)
	{
		mFrameEvents[i] = CreateEvent(nullptr, false, 0, nullptr);
		assert(mFrameEvents[i] != nullptr);
		SetEvent(mFrameEvents[i]);
	}
}

void RenderState::WaitForGame()
{
	// wait current frame until render unlock
	WaitForSingleObject(GGameState->GetFrameEvent(mFrameIndex), INFINITE);

	// lock next frame
	ResetEvent(mFrameEvents[mFrameIndex]);
}

void RenderState::MainLoop()
{
	mRenderTimer.Tick();

	// Run at loop first to clear resource for last frame
	GarbageCollection();

	// Begin the commandBuffer
	mImmediateCommandBuffer = mSwapChain->BeginRecordCommand();
}

void RenderState::AddWaitSemaphore(RefCountPtr<RHISemaphore> semaphore)
{
	mSwapChain->AddWaitSemaphore(semaphore);
}

void RenderState::AddSignalSemaphore(RefCountPtr<RHISemaphore> semaphore)
{
	mSwapChain->AddSignalSemaphore(semaphore);
}

void RenderState::SubmitTexture(RefCountPtr<RHITexture> texture)
{
	mIsSubmitTexture = true;

	if (mSubmitTexures[GetFrameIndex()] != texture)
	{
		mSwapChain->UpdateTexture(texture);
		mSubmitTexures[GetFrameIndex()] = texture;
	}
}

void RenderState::SubmitFrame()
{
#if USE_IMGUI
	mImguiGUI->Update();
	mSwapChain->EndRecordCommand([&]() { mImguiGUI->UpdateCommand(); });
#else
	mSwapChain->EndRecordCommand(nullptr);
#endif

	mSwapChain->SubmitFrame();

	GGameState->GetGameScene()->OnFrameEnd();
}

void RenderState::NextFrame()
{
	// unlock pre frame
	SetEvent(mFrameEvents[mFrameIndex]);
	mFrameIndex = (mFrameIndex + 1) % mFrameCount;
	mFrameAmount++;
}

void RenderState::OnResize()
{
	// need to check the windows is close or not, it will lose surface
	if (ThreadManager::Get().IsTimeToExit())
	{
		return;
	}

	assert(mSwapChain != nullptr);

	mSwapChain->OnResize();

	mFrameCount = mSwapChain->GetImageCount();
	for (uint32 i = 0; i < mFrameCount; ++i)
	{
		mSubmitTexures[i] = mLogoTexture;
	}

	// 矩形剪刀, 绘制全屏的话用不到
	mScissor.offset = { 0, 0 };
	mScissor.extent = mSwapChain->GetSwapchainExtent();

	// 指定视窗
	mViewport.x = 0.0f;
	mViewport.y = 0.0f;
	mViewport.width = static_cast<float>(mScissor.extent.width);
	mViewport.height = static_cast<float>(mScissor.extent.height);
	mViewport.minDepth = 0.0f;
	mViewport.maxDepth = 1.0f;

	mViewportState.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
	mViewportState.viewportCount = 1;
	mViewportState.pViewports = &mViewport;
	mViewportState.scissorCount = 1;
	mViewportState.pScissors = &mScissor;

	// Call onResize for gameScene and renderProcess
	GGameState->GetGameScene()->OnResize();
}

void RenderState::Stop()
{
	for (uint32 i = 0; i < mFrameCount; ++i)
	{
		SetEvent(mFrameEvents[i]);
	}
}

void RenderState::Exit()
{
	vkQueueWaitIdle(GVulkanRHI->GetQueue(VK_QUEUE_GRAPHICS_BIT));
	vkQueueWaitIdle(GVulkanRHI->GetQueue(VK_QUEUE_COMPUTE_BIT));
	vkQueueWaitIdle(GVulkanRHI->GetQueue(VK_QUEUE_TRANSFER_BIT));

	mSubmitTexures.clear();

	mLogoTexture.SafeRelease();

	MaggotApp::Get().DeleteRenderProcess(mRenderProcessMap);
	mRenderProcessMap.clear();

#if USE_IMGUI
	mImguiGUI->Exit();
	delete mImguiGUI;
	mImguiGUI = nullptr;
#endif

	mSwapChain->Exit();
	delete mSwapChain;
	mSwapChain = nullptr;

	// GC the last frame resource
	GarbageCollection();

	vkDestroyCommandPool(GVulkanRHI->GetLogicalDevice(), mGraphicsCommandPool, nullptr);
	vkDestroyCommandPool(GVulkanRHI->GetLogicalDevice(), mComputeCommandPool, nullptr);
	vkDestroyCommandPool(GVulkanRHI->GetLogicalDevice(), mTransferCommandPool, nullptr);
}

VkCommandBuffer RenderState::BeginSingleCommand(VkQueueFlagBits queueFlags /*= VK_QUEUE_GRAPHICS_BIT*/)
{
	// 分配一个临时命令缓冲器
	VkCommandBufferAllocateInfo allocInfo{};
	allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
	allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
	allocInfo.commandPool = GetCommandPool(queueFlags);
	allocInfo.commandBufferCount = 1;

	VkCommandBuffer commandBuffer;
	vkAllocateCommandBuffers(GVulkanRHI->GetLogicalDevice(), &allocInfo, &commandBuffer);

	// 开始记录命令缓冲区
	VkCommandBufferBeginInfo beginInfo{};
	beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
	// 表示该 Command-Buffer 只使用提交一次，用完之后就会被重置，并且每次提交时都需要重新记录
	beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;

	vkBeginCommandBuffer(commandBuffer, &beginInfo);

	return commandBuffer;
}

void RenderState::FlushSingleCommand(VkCommandBuffer commandBuffer, VkQueueFlagBits queueFlags /*= VK_QUEUE_GRAPHICS_BIT*/)
{
	VK_CHECK_RESULT(vkEndCommandBuffer(commandBuffer));

	VkSubmitInfo submitInfo{};
	submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
	submitInfo.commandBufferCount = 1;
	submitInfo.pCommandBuffers = &commandBuffer;

	VkFenceCreateInfo fenceCreateInfo{};
	fenceCreateInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
	// no need, if set VK_FENCE_CREATE_SIGNALED_BIT, the state will be signaled
	// fenceCreateInfo.flags = VK_FENCE_CREATE_SIGNALED_BIT;
	VkFence fence;
	VK_CHECK_RESULT(vkCreateFence(GVulkanRHI->GetLogicalDevice(), &fenceCreateInfo, nullptr, &fence));

	const VkQueue queue = GVulkanRHI->GetQueue(queueFlags);
	// Submit to the queue, the fence is unsignaled until finish
	VK_CHECK_RESULT(vkQueueSubmit(queue, 1, &submitInfo, fence));
	// Wait for the fence to signal that command buffer has finished executing
	VK_CHECK_RESULT(vkWaitForFences(GVulkanRHI->GetLogicalDevice(), 1, &fence, VK_TRUE, UINT64_MAX));
	vkDestroyFence(GVulkanRHI->GetLogicalDevice(), fence, nullptr);

	// vkQueueWaitIdle(queue);

	vkFreeCommandBuffers(GVulkanRHI->GetLogicalDevice(), GetCommandPool(queueFlags), 1, &commandBuffer);
}

VkCommandPool RenderState::GetCommandPool(VkQueueFlagBits queueFlags /*= VK_QUEUE_GRAPHICS_BIT*/) const
{
	switch (queueFlags)
	{
	case VK_QUEUE_GRAPHICS_BIT:
		return mGraphicsCommandPool;
		break;
	case VK_QUEUE_COMPUTE_BIT:
		return mComputeCommandPool;
		break;
	case VK_QUEUE_TRANSFER_BIT:
		return mTransferCommandPool;
		break;
	case VK_QUEUE_SPARSE_BINDING_BIT:
	case VK_QUEUE_PROTECTED_BIT:
	case VK_QUEUE_FLAG_BITS_MAX_ENUM:
	default:
		assert(0);
		break;
	}
	return mGraphicsCommandPool;
}

VkCommandBuffer RenderState::GetImmediateCommandBuffer() const
{
	return mImmediateCommandBuffer;
}

RefCountPtr<RHISampler> RenderState::GetSamplerByHash(uint32 hash) const
{
	if (mSamplerMap.find(hash) != mSamplerMap.end())
		return mSamplerMap.find(hash)->second;
	return RefCountPtr<RHISampler>();
}

void RenderState::RegisterSampler(uint32 hash, RefCountPtr<RHISampler> sampler)
{
	mSamplerMap.insert({ hash, sampler });
}

RefCountPtr<ColorBlendAttachment> RenderState::GetColorBlendAttachmentByHash(uint32 hash) const
{
	if (mColorBlendAttachmentMap.find(hash) != mColorBlendAttachmentMap.end())
		return mColorBlendAttachmentMap.find(hash)->second;
	return RefCountPtr<ColorBlendAttachment>();
}

void RenderState::RegisterColorBlendAttachment(uint32 hash, RefCountPtr<ColorBlendAttachment> colorBlendAttachment)
{
	mColorBlendAttachmentMap.insert({ hash, colorBlendAttachment });
}

uint32 RenderState::GetFrameIndex() const
{
	return mFrameIndex;
}

uint32 RenderState::GetFrameCount() const
{
	return mFrameCount;
}

uint64 RenderState::GetFrameAmount() const
{
	return mFrameAmount;
}

HANDLE RenderState::GetFrameEvent(uint32 index)
{
	return mFrameEvents[index];
}

uint32 RenderState::GetWidth() const
{
	return mScissor.extent.width;
}

uint32 RenderState::GetHeight() const
{
	return mScissor.extent.height;
}

VkFormat RenderState::GetSwapChainFormat() const
{
	assert(mSwapChain != nullptr);
	return mSwapChain->GetSwapChainFormat();
}

VkViewport& RenderState::GetViewport()
{
	return mViewport;
}

VkRect2D& RenderState::GetScissor()
{
	return mScissor;
}

VkPipelineViewportStateCreateInfo RenderState::GetViewportState() const
{
	return mViewportState;
}

float RenderState::GetTotalTime() const
{
	return mRenderTimer.GetTotalTime();
}

float RenderState::GetDeltaTime() const
{
	return mRenderTimer.GetDeltaTime();
}

uint32 RenderState::GetFps() const
{
	return mRenderTimer.GetFps();
}

RefCountPtr<RHITexture> RenderState::GetDefaultTexture()
{
	return mLogoTexture;
}

bool RenderState::IsSubmitTexture() const
{
	return mIsSubmitTexture;
}

void RenderState::ResetSubmitState()
{
	mIsSubmitTexture = false;
}

void RenderState::CreateStateResource()
{
	mLogoTexture = Picture2DTexture::Builder()
		.SetPath("MoggotEngine.png")
		.SetAutoMip()
		.Build();


}

void RenderState::CreateCommandPool()
{
	// Vulkan 需要显示声明指令缓冲（Command Buffers）用来记录所有绘图指令并且一并提交给 Vulkan，Vulkan 可以通过多线程的方法优化速度, 指令缓冲通过向设备的队列提交执行其所有指令

	VkCommandPoolCreateInfo commandPoolCreateInfo{};
	commandPoolCreateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
	commandPoolCreateInfo.queueFamilyIndex = GVulkanRHI->GetQueueFamily(VK_QUEUE_GRAPHICS_BIT);
	// 想要动态改变BUFFER的话需要设置这个flags
	commandPoolCreateInfo.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT; // Optional

	VK_CHECK_RESULT(vkCreateCommandPool(GVulkanRHI->GetLogicalDevice(), &commandPoolCreateInfo, nullptr, &mGraphicsCommandPool));

	commandPoolCreateInfo.queueFamilyIndex = GVulkanRHI->GetQueueFamily(VK_QUEUE_COMPUTE_BIT);
	VK_CHECK_RESULT(vkCreateCommandPool(GVulkanRHI->GetLogicalDevice(), &commandPoolCreateInfo, nullptr, &mComputeCommandPool));

	commandPoolCreateInfo.queueFamilyIndex = GVulkanRHI->GetQueueFamily(VK_QUEUE_TRANSFER_BIT);
	VK_CHECK_RESULT(vkCreateCommandPool(GVulkanRHI->GetLogicalDevice(), &commandPoolCreateInfo, nullptr, &mTransferCommandPool));
}

void RenderState::CreateSwapchain()
{
	mSwapChain = new SwapChain();
	mSwapChain->Init();

	// the frameCount is the same as imageCount
	mFrameCount = mSwapChain->GetImageCount();

	mSubmitTexures.resize(mFrameCount, mLogoTexture);

	// 矩形剪刀, 绘制全屏的话用不到
	mScissor.offset = { 0, 0 };
	mScissor.extent = mSwapChain->GetSwapchainExtent();

	// 指定视窗
	mViewport.x = 0.0f;
	mViewport.y = 0.0f;
	mViewport.width = static_cast<float>(mScissor.extent.width);
	mViewport.height = static_cast<float>(mScissor.extent.height);
	mViewport.minDepth = 0.0f;
	mViewport.maxDepth = 1.0f;

	mViewportState.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
	mViewportState.viewportCount = 1;
	mViewportState.pViewports = &mViewport;
	mViewportState.scissorCount = 1;
	mViewportState.pScissors = &mScissor;

#if USE_IMGUI
	mImguiGUI = new ImguiGUI(mSwapChain);
	mImguiGUI->Init();
#endif
}

void RenderState::GarbageCollection()
{
	VkDevice logicalDevice = GVulkanRHI->GetLogicalDevice();

	for (const auto& accelerationStructure : mGCAccelerationStructureMap)
	{
		vkDestroyAccelerationStructureKHR(logicalDevice, accelerationStructure, nullptr);
	}
	mGCAccelerationStructureMap.clear();

	for (const auto& buffer : mGCBufferMap)
	{
		vkDestroyBuffer(logicalDevice, buffer, nullptr);
	}
	mGCBufferMap.clear();

	for (const auto& bufferView : mGCBufferViewMap)
	{
		vkDestroyBufferView(logicalDevice, bufferView, nullptr);
	}
	mGCBufferViewMap.clear();

	for (const auto& imageView : mGCImageViewMap)
	{
		vkDestroyImageView(logicalDevice, imageView, nullptr);
	}
	mGCImageViewMap.clear();

	for (const auto& image : mGCImageMap)
	{
		vkDestroyImage(logicalDevice, image, nullptr);
	}
	mGCImageMap.clear();

	for (const auto& pipeline : mGCPipelineMap)
	{
		vkDestroyPipeline(logicalDevice, pipeline, nullptr);
	}
	mGCPipelineMap.clear();

	for (const auto& frameBuffer : mGCFrameBufferMap)
	{
		vkDestroyFramebuffer(logicalDevice, frameBuffer, nullptr);
	}
	mGCFrameBufferMap.clear();

	for (const auto& renderPass : mGCRenderPassMap)
	{
		vkDestroyRenderPass(logicalDevice, renderPass, nullptr);
	}
	mGCRenderPassMap.clear();

	for (const auto& pipelineLayout : mGCPipelineLayoutMap)
	{
		vkDestroyPipelineLayout(logicalDevice, pipelineLayout, nullptr);
	}
	mGCPipelineLayoutMap.clear();

	for (const auto& descriptorSetLayout : mGCDescriptorSetLayoutMap)
	{
		vkDestroyDescriptorSetLayout(logicalDevice, descriptorSetLayout, nullptr);
	}
	mGCDescriptorSetLayoutMap.clear();

	for (const auto& descriptorPool : mGCDescriptorPoolMap)
	{
		vkDestroyDescriptorPool(logicalDevice, descriptorPool, nullptr);
	}
	mGCDescriptorPoolMap.clear();

	for (const auto& deviceMemory : mGCDeviceMemoryMap)
	{
		vkFreeMemory(logicalDevice, deviceMemory, nullptr);
	}
	mGCDeviceMemoryMap.clear();

	for (const auto& commandBuffer : mGCCommandBufferMap)
	{
		vkFreeCommandBuffers(logicalDevice, GetCommandPool(commandBuffer.first), commandBuffer.second.size(), commandBuffer.second.data());
	}
	mGCCommandBufferMap.clear();

	for (const auto& fence : mGCFenceMap)
	{
		vkDestroyFence(logicalDevice, fence, nullptr);
	}
	mGCFenceMap.clear();

	for (const auto& semaphore : mGCSemaphoreMap)
	{
		vkDestroySemaphore(logicalDevice, semaphore, nullptr);
	}
	mGCSemaphoreMap.clear();
}

void RenderState::DelayGCBuffer(VkBuffer buffer)
{
	mGCBufferMap.push_back(buffer);
}

void RenderState::DelayGCBufferView(VkBufferView bufferView)
{
	mGCBufferViewMap.push_back(bufferView);
}

void RenderState::DelayGCDeviceMemory(VkDeviceMemory deviceMemory)
{
	mGCDeviceMemoryMap.push_back(deviceMemory);
}

void RenderState::DelayGCCommandBuffer(VkQueueFlagBits queueFlag, std::vector<VkCommandBuffer> commandBuffers)
{
	mGCCommandBufferMap.push_back({ queueFlag, commandBuffers });
}

void RenderState::DelayGCFence(VkFence fence)
{
	mGCFenceMap.push_back(fence);
}

void RenderState::DelayGCSemaphore(VkSemaphore semaphore)
{
	mGCSemaphoreMap.push_back(semaphore);
}

void RenderState::DelayGCImage(VkImage image)
{
	mGCImageMap.push_back(image);
}

void RenderState::DelayGCImageView(VkImageView imageView)
{
	mGCImageViewMap.push_back(imageView);
}

void RenderState::DelayGCPipeline(VkPipeline pipeline)
{
	mGCPipelineMap.push_back(pipeline);
}

void RenderState::DelayGCRenderPass(VkRenderPass renderPass)
{
	mGCRenderPassMap.push_back(renderPass);
}

void RenderState::DelayGCPipelineLayout(VkPipelineLayout pipelineLayout)
{
	mGCPipelineLayoutMap.push_back(pipelineLayout);
}

void RenderState::DelayGCDescriptorSetLayout(VkDescriptorSetLayout descriptorSetLayout)
{
	mGCDescriptorSetLayoutMap.push_back(descriptorSetLayout);
}

void RenderState::DelayGCDescriptorPool(VkDescriptorPool descriptorPool)
{
	mGCDescriptorPoolMap.push_back(descriptorPool);
}

void RenderState::DelayGCFrameBuffer(VkFramebuffer frameBuffer)
{
	mGCFrameBufferMap.push_back(frameBuffer);
}

void RenderState::DelayGCAccelerationStructure(VkAccelerationStructureKHR accelerationStructure)
{
	mGCAccelerationStructureMap.push_back(accelerationStructure);
}

